import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk
from PIL import Image, ImageTk
import pandas as pd
import requests
import os
import io
import threading
import queue
from concurrent.futures import ThreadPoolExecutor

class ImageViewerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图片查看器")
        self.root.geometry("600x600")
        self.root.configure(bg="#f0f0f0")

        # 初始化变量
        self.df = None
        self.current_index = -1
        self.filename = ""
        self.download_count = {}
        self.image_queue = queue.PriorityQueue()  # 使用优先级队列
        self.downloaded_indices = set()  # 跟踪已下载的索引
        self.downloaded_urls = {}  # 跟踪已下载的图片URL及其缓存路径

        # 创建 cache 和 saved_images 目录
        self.cache_dir = "cache"
        self.save_dir = "saved_images"
        for directory in [self.cache_dir, self.save_dir]:
            if not os.path.exists(directory):
                os.makedirs(directory)

        # 设置 ttk 主题
        style =(ttk.Style())
        style.theme_use("clam")
        style.configure("TButton", padding=6, font=("Arial", 10))
        style.configure("TLabel", font=("Arial", 12), background="#f0f0f0")
        style.configure("TFrame", background="#f0f0f0")

        # GUI 元素
        self.label_filename = ttk.Label(root, text="文件名: ")
        self.label_filename.pack(pady=10)

        self.label_keyword = ttk.Label(root, text="关键词: ")
        self.label_keyword.pack(pady=10)

        self.image_label = tk.Label(root, bg="#f0f0f0")
        self.image_label.pack(pady=20)

        button_frame = ttk.Frame(root)
        button_frame.pack(pady=10)

        self.previous_button = ttk.Button(button_frame, text="上一个", command=self.show_previous)
        self.previous_button.pack(side=tk.LEFT, padx=10)

        self.download_button = ttk.Button(button_frame, text="下载", command=self.save_image)
        self.download_button.pack(side=tk.LEFT, padx=10)

        self.next_button = ttk.Button(button_frame, text="下一个", command=self.show_next)
        self.next_button.pack(side=tk.LEFT, padx=10)

        # 启动图片下载线程池
        self.executor = ThreadPoolExecutor(max_workers=5)  # 5 个并行下载线程
        self.download_thread = threading.Thread(target=self.download_images, daemon=True)
        self.download_thread.start()

        # 加载 Excel 文件
        self.load_excel()

    def load_excel(self):
        file_path = filedialog.askopenfilename(filetypes=[("Excel files", "*.xlsx *.xls")])
        if not file_path:
            messagebox.showerror("错误", "未选择文件！", parent=self.root)
            self.root.quit()
            return
        try:
            self.df = pd.read_excel(file_path)
            self.show_next()
        except Exception as e:
            messagebox.showerror("错误", f"无法读取 Excel 文件: {e}", parent=self.root)
            self.root.quit()

    def download_image_task(self, index, image_url, filename):
        """单个图片下载任务"""
        # 检查是否已下载过此 URL
        if image_url in self.downloaded_urls:
            return
        cache_path = os.path.join(self.cache_dir, f"{filename}_{index}.jpg")
        # 检查缓存文件是否已存在
        if os.path.exists(cache_path):
            self.downloaded_indices.add(index)
            self.downloaded_urls[image_url] = cache_path
            return
        try:
            response = requests.get(image_url, timeout=5)
            response.raise_for_status()
            with open(cache_path, 'wb') as f:
                f.write(response.content)
            self.downloaded_indices.add(index)
            self.downloaded_urls[image_url] = cache_path
        except Exception as e:
            print(f"下载图片失败 (index {index}): {e}")

    def download_images(self):
        """从优先级队列获取任务并分配给线程池"""
        while True:
            try:
                priority, index = self.image_queue.get(timeout=1)
                if index >= len(self.df) or index < 0:
                    self.image_queue.task_done()
                    continue
                row = self.df.iloc[index]
                image_url = row['图片链接']
                filename = row['文件名']
                # 提交下载任务到线程池
                self.executor.submit(self.download_image_task, index, image_url, filename)
                self.image_queue.task_done()
            except queue.Empty:
                continue
            except Exception as e:
                print(f"队列处理错误: {e}")

    def show_next(self):
        if self.df is None:
            return
        self.current_index += 1
        if self.current_index >= len(self.df):
            messagebox.showinfo("提示", "已完成所有筛选！", parent=self.root)
            self.current_index = len(self.df) - 1
            return

        # 预先将接下来 5 张图片加入队列（优先级基于索引）
        for i in range(self.current_index, min(self.current_index + 6, len(self.df))):
            if i not in self.downloaded_indices:
                self.image_queue.put((i, i))  # (优先级, 索引)

        # 获取当前行数据
        row = self.df.iloc[self.current_index]
        filename = row['文件名']
        keyword = row['关键词']
        image_url = row['图片链接']

        # 更新文件名和关键词
        self.filename = filename
        self.label_filename.config(text=f"文件名: {filename}")
        self.label_keyword.config(text=f"关键词: {keyword}")

        # 从 cache 加载图片
        cache_path = self.downloaded_urls.get(image_url, os.path.join(self.cache_dir, f"{filename}_{self.current_index}.jpg"))
        try:
            if os.path.exists(cache_path):
                image = Image.open(cache_path)
                image = image.resize((300, 300), Image.Resampling.LANCZOS)
                self.photo = ImageTk.PhotoImage(image)
                self.image_label.config(image=self.photo)
            else:
                self.image_label.config(image="")
                messagebox.showwarning("警告", "图片尚未缓存，请稍后再试", parent=self.root)
        except Exception as e:
            self.image_label.config(image="")
            messagebox.showwarning("警告", f"无法加载图片: {e}", parent=self.root)

    def show_previous(self):
        if self.df is None:
            return
        self.current_index -= 1
        if self.current_index < 0:
            messagebox.showinfo("提示", "已经是第一张图片！", parent=self.root)
            self.current_index = 0
            return

        # 预先将前 5 张图片加入队列（优先级基于索引）
        for i in range(max(0, self.current_index - 4), self.current_index + 1):
            if i not in self.downloaded_indices:
                self.image_queue.put((i, i))  # (优先级, 索引)

        # 获取当前行数据
        row = self.df.iloc[self.current_index]
        filename = row['文件名']
        keyword = row['关键词']
        image_url = row['图片链接']

        # 更新文件名和关键词
        self.filename = filename
        self.label_filename.config(text=f"文件名: {filename}")
        self.label_keyword.config(text=f"关键词: {keyword}")

        # 从 cache 加载图片
        cache_path = self.downloaded_urls.get(image_url, os.path.join(self.cache_dir, f"{filename}_{self.current_index}.jpg"))
        try:
            if os.path.exists(cache_path):
                image = Image.open(cache_path)
                image = image.resize((300, 300), Image.Resampling.LANCZOS)
                self.photo = ImageTk.PhotoImage(image)
                self.image_label.config(image=self.photo)
            else:
                self.image_label.config(image="")
                messagebox.showwarning("警告", "图片尚未缓存，请稍后再试", parent=self.root)
        except Exception as e:
            self.image_label.config(image="")
            messagebox.showwarning("警告", f"无法加载图片: {e}", parent=self.root)

    def save_image(self):
        if self.current_index < 0 or self.df is None:
            return
        row = self.df.iloc[self.current_index]
        filename = row['文件名']
        image_url = row['图片链接']

        # 计算保存序号
        if filename not in self.download_count:
            self.download_count[filename] = 1
        else:
            self.download_count[filename] += 1
        save_index = self.download_count[filename]

        # 从 cache 复制图片到 saved_images
        cache_path = self.downloaded_urls.get(image_url, os.path.join(self.cache_dir, f"{filename}_{self.current_index}.jpg"))
        save_path = os.path.join(self.save_dir, f"{filename}_{save_index}.jpg")
        try:
            if os.path.exists(cache_path):
                with open(cache_path, 'rb') as src, open(save_path, 'wb') as dst:
                    dst.write(src.read())
                messagebox.showinfo("成功", f"图片已保存到 {save_path}", parent=self.root)
            else:
                messagebox.showerror("错误", "缓存中无此图片", parent=self.root)
        except Exception as e:
            messagebox.showerror("错误", f"保存失败: {e}", parent=self.root)

    def __del__(self):
        """清理线程池"""
        self.executor.shutdown(wait=False)

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageViewerApp(root)
    root.mainloop()